// This file implements compilation units for the Tvix compiler.
//
// In our terminology, a compilation unit is any piece of Nix code that gets
// compiled into one [bytecode.Chunk]. Functions and thunks are created from
// compilation units.
//
// Each unit has attached data which is used during compilation. The main pieces
// of data are the chunk into which bytecode is being compiled, and the scope
// tracking metadata of the compiler.
//
// Scope tracking makes up most of the complexity in this file due to Nix's
// fairly complicated scoping (with features such as mutually recursive
// bindings, `with`, frequent closure usage, ...).
//
// Tvix attempts to do as much of the heavy lifting for scoping at compile time,
// leaving the runtime to mostly deal with known stack indices. To do this, the
// compiler simulates where locals will be at runtime using the data structures
// and logic implemented for compilation units.
package compiler

import (
	"tvl.fyi/tvix/eval/internal/bytecode"
	"tvl.fyi/tvix/eval/internal/index"
)

// compilationUnit holds data (such as scope tracking) for the current chunk
// (i.e. lambda or thunk) being compiled.
type compilationUnit struct {
	chunk  bytecode.Chunk
	parent *compilationUnit

	// local variables known in this unit's scope.
	locals []local

	// byName is an index over locals. Each local has a *slice* of indices, as
	// it might be shadowed.
	byName map[string][]index.Local

	// captures represents the positions of values that need to be captured from
	// enclosing compilation unit's scopes.
	captures []index.Position

	// scopeDepth tracks how deeply the scope of this unit is nested.
	scopeDepth int
}

// Scoping logic

// local represents a single local already known to the compiler
// TODO: explicit phantoms, otherwise accessing `""` is possible and undefined.
type local struct {
	// Identifier of this local. This is always a statically known value, or a
	// "phantom" value not accessible by users.
	name string

	// Scope depth of this local
	depth int

	// Is this local already initialised? (i.e. on the stack)
	initialised bool

	// Positions of values that are captured by the thunk or closure in this slot.
	captures []index.Position

	// Is this local known to have been used at all?
	used bool
}

// stackIndex returns the runtime stack index for a given local by
// accounting for uninitialised variables at scopes below.
func (u *compilationUnit) stackIndex(target index.Local) index.Stack {
	uninitialisedCount := 0
	depth := u.locals[target].depth

	for _, l := range u.locals[:target] {
		if !l.initialised && depth > l.depth {
			uninitialisedCount++
		}
	}

	return index.Stack(int(target) - uninitialisedCount)
}

// beginScope increases the current scope depth (e.g. for a new
// bindings block, or `with`-scope).
func (u *compilationUnit) beginScope() {
	u.scopeDepth++
}

// endScope closes a scope and removes all of its associated still
// tracked locals.
//
// Returns the count of locals that were dropped while marked as
// initialised, which is used to determine whether to emit scope
// cleanup operations.
func (u *compilationUnit) endScope() int {
	if u.scopeDepth == 0 {
		panic("tvix bug: can not end top scope in compiler")
	}

	pops := 0

	// TL;DR - iterate from the back while things belonging to the
	// ended scope still exist.
	for i := len(u.locals) - 1; i >= 0; i-- {
		l := u.locals[i]
		if l.depth != u.scopeDepth {
			break
		}

		if l.initialised {
			pops++
		}

		// TODO: emit positions of unused locals here

		// remove by-name index if this was a named local
		if l.name != "" {
			indices := u.byName[l.name]
			u.byName[l.name] = indices[:len(indices)-1]
		}

		u.locals = u.locals[:len(u.locals)-1]
		if len(u.locals) == 0 {
			break
		}
	}

	u.scopeDepth--
	return pops
}

const notShadowed index.Local = -1

// declareLocal declares an unitialised local variable, and returns
// its new local index and optionally the index of a previous local if
// it was shadowed.
func (u *compilationUnit) declareLocal(name string) (index.Local, index.Local) {
	idx := index.Local(len(u.locals))
	u.locals = append(u.locals, local{
		name:        name,
		depth:       u.scopeDepth,
		initialised: false,
		used:        false,
	})

	if u.byName == nil {
		u.byName = make(map[string][]index.Local)
	} else if shadows, _ := u.byName[name]; len(shadows) > 0 {
		u.byName[name] = append(shadows, idx)
		return idx, shadows[0]
	}

	u.byName[name] = []index.Local{idx}
	return idx, notShadowed
}

// declareLocal declares a local variable that occupies a stack slot,
// but is not directly accessible by users (e.g. attribute sets used
// for `with`).
func (u *compilationUnit) declarePhantom() index.Local {
	idx := index.Local(len(u.locals))
	u.locals = append(u.locals, local{
		depth:       u.scopeDepth,
		initialised: false,
		used:        true, // phantoms are always used by some instruction
	})

	return idx
}

// resolveLocal looks up a variable by name in the compilation unit's scope, and
// recurses into the parent scopes to potentially capture an enclosing variable.
func (u *compilationUnit) resolveLocal(name string) index.Position {
	pos := index.Unknown

	if indices, ok := u.byName[name]; ok {
		idx := index.Local(indices[len(indices)-1])
		u.locals[idx].used = true
		pos = index.NewLocalPosition(idx)
		return pos
	}

	if u.parent != nil {
		nested := u.parent.resolveLocal(name)
		if nested != index.Unknown {
			pos = index.NewCapturePosition(u.capture(pos))
		}
	}

	return pos
}

func (u *compilationUnit) capture(pos index.Position) index.Capture {
	// Check if this value has already been captured.
	for i, p := range u.captures {
		if p == pos {
			return index.Capture(i)
		}
	}

	idx := index.Capture(len(u.captures))
	u.captures = append(u.captures, pos)
	return idx
}
